import argparse
import torch.multiprocessing as mp
import torch
import os
import sys
import glob

if os.getcwd() not in sys.path:
    sys.path.append(os.getcwd())

from utils.utils import setup_seed
from utils.utils import get_source_list
from utils.utils import list_to_files

from wavlm.WavLMWrapper import WavLMWrapper as WavLM
import torchaudio
import tqdm
import numpy as np

SEED = 1234


def inference(rank: int, args: argparse.Namespace):
    device = args.gpus[rank % len(args.gpus)]
    source_list = get_source_list(args.audio_scp)[
        rank :: args.num_proc
    ]  # list of audio paths
    wavlm = WavLM(args.wavlm_ckpt).to(device)
    print(f"rank {rank} got data number {len(source_list)}")
    print(f"output directory {args.output_dir}")
    wavlm_emb_path = os.path.join(args.output_dir, f"wavlm_l{args.layer_num}", "audio")
    scp_path = os.path.join(
        args.output_dir, f"wavlm_l{args.layer_num}", f"audio_{rank}.scp"
    )
    length_path = os.path.join(
        args.output_dir, f"wavlm_l{args.layer_num}", f"audio_len_{rank}.scp"
    )
    os.makedirs(wavlm_emb_path, exist_ok=True)
    result = []
    length = []
    with torch.no_grad():
        for s in tqdm.tqdm(source_list):
            s: str
            audio, _ = torchaudio.load(s)
            audio = audio.to(device)  # [1,T]
            embedding = wavlm(audio).squeeze(0).cpu().numpy()  # np [T',E]
            output_path = (
                os.path.splitext(os.path.join(wavlm_emb_path, s.split("/")[-1]))[0]
                + ".npy"
            )
            t = embedding.shape[0]
            np.save(output_path, embedding)
            result.append(output_path)
            length.append(t)
            pass
    print("writing list to files")
    list_to_files(result, scp_path)
    result_length = []
    for r, l in zip(result, length):
        result_length.append(f"{r} {l}")
        pass
    list_to_files(result_length, length_path)
    

def main(args):
    # os.makedirs(args.)
    setup_seed(SEED)
    mp.spawn(inference, args=(args,), nprocs=args.num_proc, join=True)
    print("inference done")
    length_path = os.path.join(
        args.output_dir, f"wavlm_l{args.layer_num}"
    )
    audio_scps = glob.glob(f"{length_path}/audio_[0-9]*.scp")
    audio_len_scps = glob.glob(f"{length_path}/audio_len_[0-9]*.scp")
    print("saving the best audios")
    
    res = []
    for p in audio_scps:
        li = get_source_list(p)
        res = res + li
    list_to_files(res, f"{length_path}/audio.scp")

    res = []
    for p in audio_len_scps:
        li = get_source_list(p)
        res = res + li
    list_to_files(res, f"{length_path}/audio_len.scp")

    pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--audio_scp",
        type=str,
        required=True,
        help="""
            The file has a list of source paths, where each path is on each line. It looks like:
            ...
            >>> path/to/wav/1.wav
            >>> path/to/wav/2.wav
            ...
            """,
    )
    parser.add_argument("--output_dir", type=str, required=True)
    parser.add_argument("--wavlm_ckpt", type=str, default="./ckpt/WavLM-Large.pt")
    parser.add_argument("--layer_num", type=int, default=6)
    parser.add_argument(
        "--num_proc", type=int, default=8, help="total number of procedures"
    )
    parser.add_argument(
        "--gpus", nargs="+", default=["cuda:0", "cuda:1", "cuda:2", "cuda:3"]
    )
    args = parser.parse_args()
    main(args)
    pass
